// vm.js

// 指令集 (OPCODES) - 重新分配操作码，确保NOP为0，并将算术/位运算指令连续排列
const OPCODES = {
    NOP: 0,              // 无操作
    PUSH: 1,             // 压入栈
    POP: 2,              // 弹出栈顶元素

    // 算术运算 (连续)
    ADD: 3,              // 加法
    SUB: 4,              // 减法
    MUL: 5,              // 乘法
    DIV: 6,              // 除法
    MOD: 7,              // 取余

    // 位运算 (连续)
    AND: 8,              // 按位与
    OR: 9,               // 按位或
    NOT: 10,             // 按位非

    // 比较运算
    EQ: 11,              // 等于
    GT: 12,              // 大于
    LT: 13,              // 小于

    // 内存操作
    LOAD: 14,            // 从指定地址加载
    STORE: 15,           // 存储到指定地址
    STORE_INDIRECT: 20,  // 间接存储
    LOAD_INDIRECT: 21,   // 间接加载

    // 控制流
    JUMP: 16,            // 无条件跳转
    JUMPIF: 17,          // 条件跳转

    // I/O & 程序终止
    PRINT: 18,           // 打印栈顶
    HALT: 19             // 停止
};

// 指令的参数数量 (用于PC步进) - 映射到新的操作码
const OP_ARG_COUNTS = {
    [OPCODES.NOP]: 0,
    [OPCODES.PUSH]: 1,
    [OPCODES.POP]: 0,
    [OPCODES.ADD]: 0,
    [OPCODES.SUB]: 0,
    [OPCODES.MUL]: 0,
    [OPCODES.DIV]: 0,
    [OPCODES.MOD]: 0,
    [OPCODES.AND]: 0,
    [OPCODES.OR]: 0,
    [OPCODES.NOT]: 0,
    [OPCODES.EQ]: 0,
    [OPCODES.GT]: 0,
    [OPCODES.LT]: 0,
    [OPCODES.LOAD]: 1,
    [OPCODES.STORE]: 1,
    [OPCODES.JUMP]: 1,
    [OPCODES.JUMPIF]: 1,
    [OPCODES.PRINT]: 0,
    [OPCODES.HALT]: 0,
    [OPCODES.STORE_INDIRECT]: 0,
    [OPCODES.LOAD_INDIRECT]: 0
};

/**
 * StackVM 类：基于栈的虚拟机核心
 */
class StackVM {
    constructor(memorySize = 256) {
        this.memory = new Int32Array(memorySize);
        this.stack = [];
        this.program = [];
        this.pc = 0;
        this.isRunning = false;
        this.outputBuffer = [];
        this.cycleCount = 0; // 用于记录指令周期

        this.lastMemWriteAddr = -1;
        this.lastExecutedPc = -1;

        this.OPCODES = OPCODES;
        this.OP_ARG_COUNTS = OP_ARG_COUNTS;

        this.onStateChange = null;
        this.onOutput = null;
    }

    // 加载程序
    loadProgram(bytecodeArray) {
        this.program = bytecodeArray;
        this.reset();
    }

    // 重置虚拟机状态
    reset() {
        this.stack = [];
        this.memory.fill(0); // Int32Array可以使用fill来清零
        this.pc = 0;
        this.isRunning = false;
        this.outputBuffer = [];
        this.cycleCount = 0; // 重置周期计数
        this.lastMemWriteAddr = -1;
        this.lastExecutedPc = -1;
        this.onStateChange?.(this.getVMState()); // 通知UI更新
    }

    // 获取当前虚拟机状态 (用于UI显示)
    getVMState() {
        return {
            memory: this.memory, // 直接传递 TypedArray
            stack: [...this.stack], // 复制一份栈，避免外部修改
            pc: this.pc,
            currentInstruction: this.program[this.pc],
            currentOperand: this.program[this.pc + 1],
            isRunning: this.isRunning,
            program: [...this.program], // 复制一份程序，避免外部修改
            output: [...this.outputBuffer],
            lastMemWriteAddr: this.lastMemWriteAddr,
            lastExecutedPc: this.lastExecutedPc,
            cycleCount: this.cycleCount // 将周期数暴露给UI
        };
    }

    // 执行一步指令
    step() {
        if (!this.isRunning || this.pc >= this.program.length) {
            this.isRunning = false;
            this.onStateChange?.(this.getVMState());
            return false;
        }

        this.cycleCount++; // 每执行一步，周期数加一

        const initialPc = this.pc;
        const opcode = this.program[this.pc];
        const operand = this.program[this.pc + 1];
        let pcIncrement = 1; // 默认PC只前进1步（操作码本身）

        this.lastMemWriteAddr = -1; // 重置上一次内存写入地址

        try {
            switch (opcode) {
                case this.OPCODES.NOP:
                    // NOP指令，不执行任何操作，仅推进PC
                    break;
                case this.OPCODES.PUSH:
                    this.stack.push(operand);
                    pcIncrement += this.OP_ARG_COUNTS[opcode];
                    break;
                case this.OPCODES.POP:
                    if (this.stack.length < 1) throw new Error("Stack underflow: POP on empty stack.");
                    this.stack.pop();
                    break;
                // 算术和位运算指令
                case this.OPCODES.ADD:
                case this.OPCODES.SUB:
                case this.OPCODES.MUL:
                case this.OPCODES.DIV:
                case this.OPCODES.MOD:
                case this.OPCODES.AND:
                case this.OPCODES.OR:
                    if (this.stack.length < 2) throw new Error("Stack underflow for arithmetic/bitwise operation.");
                    const b = this.stack.pop();
                    const a = this.stack.pop();
                    let result;
                    switch (opcode) {
                        case this.OPCODES.ADD: result = a + b; break;
                        case this.OPCODES.SUB: result = a - b; break;
                        case this.OPCODES.MUL: result = a * b; break;
                        case this.OPCODES.DIV:
                            if (b === 0) throw new Error("Division by zero.");
                            result = Math.floor(a / b);
                            break;
                        case this.OPCODES.MOD:
                            if (b === 0) throw new Error("Modulo by zero.");
                            result = a % b;
                            break;
                        case this.OPCODES.AND: result = a & b; break;
                        case this.OPCODES.OR: result = a | b; break;
                    }
                    this.stack.push(result);
                    break;
                case this.OPCODES.NOT:
                    if (this.stack.length < 1) throw new Error("Stack underflow for NOT operation.");
                    const val_not = this.stack.pop();
                    this.stack.push(~val_not);
                    break;
                // 比较运算
                case this.OPCODES.EQ:
                case this.OPCODES.GT:
                case this.OPCODES.LT:
                    if (this.stack.length < 2) throw new Error("Stack underflow for comparison operation.");
                    const valB = this.stack.pop();
                    const valA = this.stack.pop();
                    let compResult;
                    switch (opcode) {
                        case this.OPCODES.EQ: compResult = (valA === valB) ? 1 : 0; break;
                        case this.OPCODES.GT: compResult = (valA > valB) ? 1 : 0; break;
                        case this.OPCODES.LT: compResult = (valA < valB) ? 1 : 0; break;
                    }
                    this.stack.push(compResult);
                    break;
                // 内存操作
                case this.OPCODES.LOAD:
                    if (operand === undefined || operand < 0 || operand >= this.memory.length) {
                        throw new Error(`Invalid memory address for LOAD: ${operand}`);
                    }
                    this.stack.push(this.memory[operand]);
                    pcIncrement += this.OP_ARG_COUNTS[opcode];
                    break;
                case this.OPCODES.STORE:
                    if (this.stack.length < 1) throw new Error("Stack underflow for STORE.");
                    if (operand === undefined || operand < 0 || operand >= this.memory.length) {
                        throw new Error(`Invalid memory address for STORE: ${operand}`);
                    }
                    this.memory[operand] = this.stack.pop();
                    this.lastMemWriteAddr = operand;
                    pcIncrement += this.OP_ARG_COUNTS[opcode];
                    break;
                case this.OPCODES.LOAD_INDIRECT:
                    if (this.stack.length < 1) throw new Error("Stack underflow for LOAD_INDIRECT: address not on stack.");
                    const addrToLoad = this.stack.pop();
                    if (typeof addrToLoad !== 'number' || addrToLoad < 0 || addrToLoad >= this.memory.length) {
                        throw new Error(`Invalid memory address for LOAD_INDIRECT: ${addrToLoad}`);
                    }
                    this.stack.push(this.memory[addrToLoad]);
                    break;
                case this.OPCODES.STORE_INDIRECT:
                    if (this.stack.length < 2) throw new Error("Stack underflow for STORE_INDIRECT: value and address not on stack.");
                    const addrToStore = this.stack.pop();
                    const valToStore = this.stack.pop();
                    if (typeof addrToStore !== 'number' || addrToStore < 0 || addrToStore >= this.memory.length) {
                        throw new Error(`Invalid memory address for STORE_INDIRECT: ${addrToStore}`);
                    }
                    this.memory[addrToStore] = valToStore;
                    this.lastMemWriteAddr = addrToStore;
                    break;
                // 控制流
                case this.OPCODES.JUMP:
                    if (operand === undefined || operand < 0 || operand >= this.program.length) {
                        throw new Error(`Invalid jump address: ${operand}`);
                    }
                    this.pc = operand;
                    this.lastExecutedPc = initialPc;
                    this.onStateChange?.(this.getVMState());
                    return true;
                case this.OPCODES.JUMPIF:
                    if (this.stack.length < 1) throw new Error("Stack underflow for JUMPIF.");
                    const condition = this.stack.pop();
                    if (condition !== 0) {
                        if (operand === undefined || operand < 0 || operand >= this.program.length) {
                            throw new Error(`Invalid jump address for JUMPIF: ${operand}`);
                        }
                        this.pc = operand;
                        this.lastExecutedPc = initialPc;
                        this.onStateChange?.(this.getVMState());
                        return true;
                    }
                    pcIncrement += this.OP_ARG_COUNTS[opcode];
                    break;
                // I/O & 程序终止
                case this.OPCODES.PRINT:
                    if (this.stack.length < 1) throw new Error("Stack underflow for PRINT.");
                    const valueToPrint = this.stack.pop();
                    this.outputBuffer.push(valueToPrint);
                    this.onOutput?.(valueToPrint);
                    break;
                case this.OPCODES.HALT:
                    this.isRunning = false;
                    break;
                default:
                    throw new Error(`Unknown opcode: ${opcode} at PC: ${this.pc}`);
            }
        } catch (e) {
            this.isRunning = false;
            this.onOutput?.(`RUNTIME ERROR at PC ${this.pc}: ${e.message}`);
            throw e;
        }

        this.pc += pcIncrement;
        this.lastExecutedPc = initialPc;

        this.onStateChange?.(this.getVMState());
        return this.isRunning;
    }

    // 运行整个程序
    run() {
        this.isRunning = true;
        this.outputBuffer = [];
        this.onStateChange?.(this.getVMState());

        const executeLoop = () => {
            if (this.isRunning && this.pc < this.program.length) {
                try {
                    const continueRunning = this.step();
                    if (continueRunning) {
                        setTimeout(executeLoop, 10);
                    } else {
                        this.onStateChange?.(this.getVMState());
                    }
                } catch (e) {
                    console.error("VM Run Error:", e.message);
                    this.isRunning = false;
                    this.onStateChange?.(this.getVMState());
                }
            } else {
                this.isRunning = false;
                this.onStateChange?.(this.getVMState());
            }
        };
        executeLoop();
    }
}


/**
 * Assembler 函数：将汇编代码转换为字节码
 */
function assemble(sourceCode) {
    const lines = sourceCode.split('\n')
        .map(line => line.split(';')[0].trim())
        .filter(line => line.length > 0);
    
    const bytecode = [];
    const labels = {};
    const unresolvedJumps = [];

    // 别名映射 (包含新的位运算别名)
    const ALIAS_MAP = {
        '+': 'ADD', '-': 'SUB', '*': 'MUL', '/': 'DIV',
        '=': 'EQ', '>': 'GT', '<': 'LT', '%': 'MOD',
        '&': 'AND', '|': 'OR', '~': 'NOT'
    };

    // 第一遍扫描：识别所有标签并预计算指令长度
    let currentAddress = 0;
    for (let i = 0; i < lines.length; i++) {
        const line = lines[i];
        if (line.endsWith(':')) {
            const labelName = line.slice(0, -1);
            if (labels[labelName] !== undefined) {
                throw new Error(`Duplicate label '${labelName}' at line ${i + 1}`);
            }
            labels[labelName] = currentAddress;
        } else {
            const parts = line.split(/\s+/);
            let mnemonic = parts[0].toUpperCase();

            if (ALIAS_MAP[mnemonic]) {
                mnemonic = ALIAS_MAP[mnemonic];
            }

            if (OPCODES[mnemonic] === undefined) {
                throw new Error(`Unknown instruction or alias '${parts[0]}' at line ${i + 1}`);
            }
            if (OP_ARG_COUNTS[OPCODES[mnemonic]] === undefined) {
                throw new Error(`Missing argument count for instruction '${mnemonic}' in OP_ARG_COUNTS at line ${i + 1}`);
            }
            currentAddress += (1 + OP_ARG_COUNTS[OPCODES[mnemonic]]);
        }
    }

    // 第二遍扫描：生成字节码并解析跳转
    for (let i = 0; i < lines.length; i++) {
        const line = lines[i];
        if (line.endsWith(':')) {
            continue;
        }

        const parts = line.split(/\s+/);
        let mnemonic = parts[0].toUpperCase();
        const operandStr = parts[1];

        if (ALIAS_MAP[mnemonic]) {
            mnemonic = ALIAS_MAP[mnemonic];
        }

        const opcode = OPCODES[mnemonic];
        bytecode.push(opcode);

        const expectedArgs = OP_ARG_COUNTS[opcode];

        if (expectedArgs > 0) {
            if (operandStr === undefined) {
                throw new Error(`Instruction '${mnemonic}' at line ${i + 1} requires an operand.`);
            }

            if (mnemonic === 'JUMP' || mnemonic === 'JUMPIF') {
                if (labels[operandStr] === undefined) {
                    unresolvedJumps.push({
                        line: i + 1,
                        bytecodeIndex: bytecode.length,
                        label: operandStr
                    });
                    bytecode.push(0); // 暂时占位
                } else {
                    bytecode.push(labels[operandStr]);
                }
            } else {
                let operand;
                if (operandStr.startsWith('0x')) {
                    operand = parseInt(operandStr.slice(2), 16);
                } else {
                    operand = parseInt(operandStr, 10);
                }
                
                if (isNaN(operand)) {
                    throw new Error(`Invalid numeric operand '${operandStr}' for '${mnemonic}' at line ${i + 1}`);
                }
                bytecode.push(operand);
            }
        } else if (operandStr !== undefined) {
            // 对不需要操作数的指令，如果提供了操作数，则报错。
            // 排除 LOAD_INDIRECT / STORE_INDIRECT 这些特殊指令，因为它们的操作数来自栈。
            // NOP 也不应该有操作数
            if (![OPCODES.LOAD_INDIRECT, OPCODES.STORE_INDIRECT, OPCODES.NOP].includes(opcode)) {
                throw new Error(`Instruction '${mnemonic}' at line ${i + 1} does not take an operand.`);
            } else if (opcode === OPCODES.NOP) {
                 throw new Error(`NOP instruction at line ${i + 1} does not take an operand.`);
            }
        }
    }

    for (const jump of unresolvedJumps) {
        if (labels[jump.label] === undefined) {
            throw new Error(`Undefined label '${jump.label}' referenced at line ${jump.line}`);
        }
        bytecode[jump.bytecodeIndex] = labels[jump.label];
    }

    return bytecode;
}